'''
    基于线程安全的 任务队列实现
'''

import threading
import time

class ThreadSafeQueueException(Exception):
    """ 定义异常 """
    pass


# 线程安全的队列
class ThreadSafeQueue():
    """
        加锁的过程，可以直接使用上下文协议（with）https://docs.python.org/zh-cn/3.9/library/threading.html#condition-objects
        with 条件变量， with lock 是一样的
    """

    def __init__(self, max_size = 0):
        """
            param:  max_size: 默认为0，表示无限大
            queue 用一个列表表示
            实现线程安全需要加锁 及 条件变量
        """
        self.queue = list()
        self.max_size = max_size
        self.lock = threading.Lock()
        self.condition = threading.Condition(self.lock)

    def size(self):
        """ 当前队列元素数量 """
        with self.lock:
            return len(self.queue)


    def put(self, item):
        """ 添加元素 """
        # 如果当前 元素已满， 不可插入
        if self.max_size != 0 and self.size() > self.max_size:
            raise ThreadSafeQueueException
        # 否则加锁: 此时with上下文管理，相当于对底层的锁，进行 加锁解锁
        with self.condition:
           self.queue.append(item) 
           # 通知其他等待 获取元素的线程，执行get 获取
           self.condition.notify()


    
    def batch_put(self, item_list: list):
        """ 批量放入 
        param:  item_list: 添加元素列表， 不是列表要转化为列表
        """
        if not isinstance(item_list, list):
            item_list = list(item_list)
        
        # 依次添加
        for item in item_list:
            self.put(item)
        

    def pop(self, block = False, timeout = None):
        """ 取元素 
            param:  block 当队列没有元素，是否阻塞的去等待
            param:  timeout 如果等待，等待多久
        """
        with self.condition:
            # 其他无法会的锁，无法执行。 这里直接len(self.queue), 可以
            if len(self.queue) == 0:
                # 如果需要阻塞等待
                if block:
                    if timeout == None:
                        timeout = 0
                    while len(self.queue) <= 0:
                        print("等待")
                        self.condition.wait(timeout = timeout)
                        print("等待结束")
                        # 继续按照官网的写法
                        # 条件变量满足，线程重新唤醒
                        # wait返回前，重新拿到互斥锁
                        # 睡眠这段时间，至少有一个生产者压入新任务
                        # 新元素 有可能被其他线程抢走，因此需要重新检查堆是否为空，要查看 notify 官网的用法
                        continue
                else:
                    return None  
            # 执行获得元素操作
            item = self.queue.pop()
            # put 有对应的逻辑的话，应该有个通知唤醒步骤
            # self.condition.notify()
            return item


    def get(self, index):
        """ 获取指定 index 的元素 """
        with self.condition:
            if index > len(self.queue) - 1:
                return None
            item = self.queue[index]
            return item


# 测试，一个生产者，两个消费者
if __name__ == '__main__':
    queue = ThreadSafeQueue(max_size = 100)
    def producer():
        while True:
            queue.put(1)
            time.sleep(3)

    def consumer1():
        while True:
            item = queue.pop(block= True, timeout= 2)
            print(item, "consumer1 拿到啦")
            time.sleep(1)
    
    def consumer2():
        while True:
            item = queue.pop(block= True, timeout= 2)
            print(item, "consumer2 拿到了")
            time.sleep(1)

    thread1 = threading.Thread(target= producer)
    thread2 = threading.Thread(target= consumer1)
    thread3 = threading.Thread(target= consumer2)
    thread1.start()
    thread2.start()
    thread3.start()
    thread1.join()
    thread2.join()
    thread3.join()


    
